'''
Created on Jun 15, 2011

@author: oabalbin
'''

from optparse import OptionParser
from exome.jobs.base import JOB_SUCCESS, JOB_ERROR
from exome.jobs.job_runner import qsub_cac, qsub_loc, run_local
from exome.jobs.config import ExomePipelineConfig, ExomeAnalysisConfig
from exome.jobs.cohorts import define_patients_cohort, define_sample_cohort
from collections import defaultdict, deque


# Global Variables
NODE_MEM=45000.0
NODE_CORES=12
SINGLE_CORE=1
MEM_PER_CORE= int(float(NODE_MEM) / NODE_CORES)
# wt=walltime
WT_SHORT= "24:00:00"
WT_LONG= "100:00:00"#"60:00:00" "100:00:00"


def samtools_mpileup(ref_genome, list_bam_files, outfile, path_to_sam, target_regions=None):
    '''
    Returns the shell to execute a pile up using sam tools 
    param is a dictionary with the input param
    '''
    
    samtools_command=path_to_sam+'samtools'
    if target_regions is not None:
        args=[samtools_command,'mpileup','-f',ref_genome] + list_bam_files +['-l', target_regions, '>', outfile]
    else:
        args=[samtools_command,'mpileup','-f',ref_genome] + list_bam_files +['>', outfile]
        
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command


def samtools_bcf_mpileup(ref_genome, list_bam_files, outfile, path_to_sam, target_regions=None):
    '''
    Returns the shell to execute a pile up using sam tools 
    param is a dictionary with the input param
    '''
    
    samtools_command=path_to_sam+'samtools'
    if target_regions is not None:
        args=[samtools_command,'mpileup','-gf',ref_genome] + list_bam_files +['-l', target_regions, '>', outfile]
    else:
        args=[samtools_command,'mpileup','-gf',ref_genome] + list_bam_files +['>', outfile]
        
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command


def bcftools_call_snps(bcf_file_name, outfile_name, path_to_sam):
    '''
    Call snps candidates using bcftools
    This tool converts BCF to VCF, indexes BCF for random access, concatenates (not merges) BCFs, 
    estimate site allele frequencies and calls SNP candidates.
    '''
    bcftools_command=path_to_sam+'bcftools/bcftools'
    args = [bcftools_command, 'view', '-vcg', bcf_file_name, '>', outfile_name]
    
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command


def sam_bcftools_mpileup(ref_genome, list_bam_files, outfile, path_to_sam, target_regions=None):
    '''
    '''
    samtools_command=path_to_sam+'samtools'
    bcftools_command=path_to_sam+'bcftools/bcftools'
    outtem=open(outfile,'w')
    outtem.close()
    # 06-17-11
    # pileup is used instead of mpileup. Because in varscan and other softwares
    # are not adapted to mpileup. 
    # Note that pileup is deprecated, and would need to be replaced when using 
    # more recent versions of samtools (currently used v1.10)
    if target_regions is not None:
        args=[samtools_command,'pileup','-ugf',ref_genome] + list_bam_files +['-l', target_regions, '|'] +\
        [bcftools_command, 'view', '-vcg', '-', '>', outfile]
    else:
        args=[samtools_command,'pileup','-ugf',ref_genome] + list_bam_files +['|']+\
        [bcftools_command, 'view', '-vcg', '-', '>', outfile]
        
    args= [a.replace(',',';') for a in args]
    command = ",".join(args).replace(',',' ').replace(';',',')

    return command



def do_pileup(analysis, configrun, use_recal_files, jobrunfunc, prev_deps, analysis_type='cohort'):
    '''
    It performs pileups using samtools.
    This pileups are used for: calculating coverage
    Call SNVs using samtools, varscan, others.
    '''
    
    extra_mem, num_cores = configrun.gatk_use_mem, configrun.gatk_num_cores
    path_to_picard, path_to_sam = configrun.picard_path, configrun.samtools_path
    target_exons = configrun.sam_target_exons
    genomes = configrun.genomes['human']
    ref_genome = genomes.gatk_ref_genome    
    my_email=configrun.email_addresses

    pileup_deps=[]
    list_bam_samples=deque()
    if analysis_type=='bypatient':
        for sp in analysis.samples:
            jobn='pup_'+sp.name+sp.category
            if sp.category=='benign':
                if use_recal_files:
                    list_bam_samples.append(sp.sorted_mmarkdup_bam)
                else:
                    list_bam_samples.append(sp.sorted_quickmmdup_bam)
                    
                bcf_file_name=sp.benign_bcf_mpileup_file
                pileup_file_name=sp.benign_reg_mpileup_file
            else:
                if use_recal_files:
                    list_bam_samples.append(sp.sorted_mmarkdup_bam)
                else:
                    list_bam_samples.append(sp.sorted_quickmmdup_bam)

                bcf_file_name=sp.tumor_bcf_mpileup_file
                pileup_file_name=sp.tumor_reg_mpileup_file        
                                
                '''
                cohort=define_patients_cohort(analysis,use_recal_files)
                
                for pid, cohort_samples in cohort.iteritems():
                    for sptype, list_bam_samples in cohort_samples.iteritems():
                        jobn='pup_'+pid+sptype
                        
                        if sptype == 'benign':
                            bcf_file_name=analysis.benign_bcf_mpileup_file
                            pileup_file_name=analysis.benign_mpileup_file
                        elif sptype == 'tumor':
                            bcf_file_name=analysis.tumor_bcf_mpileup_file
                            pileup_file_name=analysis.tumor_mpileup_file
                '''      
                '''
                command = samtools_bcf_mpileup(ref_genome, list(list_bam_samples), bcf_file_name, path_to_sam, target_exons)
                # Create the bcf mpileup file
                jobidbcf = jobrunfunc(jobn+'_bcf', command, SINGLE_CORE, cwd=None, walltime=WT_LONG, pmem=extra_mem, 
                                   deps=prev_deps, stdout=None, email_addresses=my_email)
                '''
            command = samtools_mpileup(ref_genome, list(list_bam_samples), pileup_file_name, path_to_sam, target_exons)
            print command
            # Create the bcf mpileup file
            jobidpup = jobrunfunc(jobn+'_pup', command, SINGLE_CORE, cwd=None, walltime=WT_LONG, pmem=extra_mem, 
                               deps=prev_deps, stdout=None, email_addresses=my_email)
            
            pileup_deps.append(jobidpup)
        
    elif analysis_type=='cohort':
        cohort = define_sample_cohort(analysis,use_recal_files)    
        for sptype, list_bam_samples in cohort.iteritems():
            #
            jobn='pp'
            jobn=jobn+sptype
            
            if sptype == 'benign':
                bcf_file_name=analysis.benign_bcf_mpileup_file
                pileup_file_name=analysis.benign_reg_mpileup_file
            elif sptype == 'tumor':
                bcf_file_name=analysis.tumor_bcf_mpileup_file
                pileup_file_name=analysis.tumor_reg_mpileup_file
                            
            '''
            command = samtools_bcf_mpileup(ref_genome, list(list_bam_samples), bcf_file_name, path_to_sam, target_exons)
            # Create the bcf mpileup file
            jobidbcf = jobrunfunc(jobn+'_bcf', command, SINGLE_CORE, cwd=None, walltime=WT_LONG, pmem=extra_mem, 
                               deps=prev_deps, stdout=None, email_addresses=my_email)
            '''
            command = samtools_mpileup(ref_genome, list(list_bam_samples), pileup_file_name, path_to_sam, target_exons)
            # Create the bcf mpileup file
            jobidpup = jobrunfunc(jobn+'_pup', command, SINGLE_CORE, cwd=None, walltime=WT_LONG, pmem=extra_mem, 
                               deps=prev_deps, stdout=None, email_addresses=my_email)
            
            pileup_deps.append(jobidpup)

            
    return pileup_deps
            
            
            
if __name__ == '__main__':
    
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-r", "--config_file", dest="config_file",
                            help="file with run configuration")
    optionparser.add_option("-a", "--analysis_file", dest="analysis_file",
                            help="file with experiment configuration") 
    optionparser.add_option("--use_recal_files", dest="use_recal_files", action="store_true", default=False,
                            help="use recal files to call the snvs variants")
    optionparser.add_option("-t","--analysis_type", dest="analysis_type", default='cohort',
                            help="use bypatient or cohort")

    
    optionparser.add_option("--local_cluster", dest="local_cluster", action="store_true", default=False)
    optionparser.add_option("--local", dest="local", action="store_true", default=False)
    optionparser.add_option("--cluster", dest="cluster", action="store_true", default=False)

    (options, args) = optionparser.parse_args()    

    config = ExomePipelineConfig()
    config.from_xml(options.config_file)
    analysis = ExomeAnalysisConfig()
    analysis.from_xml(options.analysis_file, config.output_dir)
    #Default when called from the command line
    
    depends=[]
    
    if not (options.local ^ options.cluster ^ options.local_cluster):
        optionparser.error("Must set either --local, --cluster or --local_cluster to run job")
    if options.local:
        jobrunfunc = run_local
    elif options.cluster:
        jobrunfunc = qsub_cac
    elif options.local_cluster:
        jobrunfunc = qsub_loc
        
    do_pileup(analysis, config, options.use_recal_files, jobrunfunc, depends, options.analysis_type)

            
